import sqlite3
import threading
import queue
from typing import Optional
from utils.logger import logger

class DatabasePool:
    def __init__(self, max_connections=5):
        self.db_path = None
        self.pool = queue.Queue(maxsize=max_connections)
        self.size = max_connections
        self.lock = threading.RLock()
        self._is_initialized = False
        
    def initialize(self, db_path):
        """初始化连接池"""
        with self.lock:
            # 如果已初始化且路径相同，直接返回
            if self._is_initialized and self.db_path == db_path:
                return True
                
            # 清空旧连接池
            self._clear_pool()
            
            # 设置新路径
            self.db_path = db_path
            
            # 创建连接并放入池中
            try:
                for _ in range(self.size):
                    conn = sqlite3.connect(db_path, check_same_thread=False)
                    # 启用外键约束、WAL模式和busy_timeout
                    conn.execute("PRAGMA foreign_keys = ON")
                    conn.execute("PRAGMA journal_mode = WAL")
                    conn.execute("PRAGMA busy_timeout = 10000")  # 10秒锁等待超时
                    conn.execute("PRAGMA synchronous = NORMAL")  # 降低同步级别，提高性能
                    self.pool.put(conn)
                    
                logger.info(f"数据库连接池初始化成功，连接数: {self.size}, 路径: {db_path}")
                self._is_initialized = True
                return True
            except Exception as e:
                logger.error(f"初始化连接池失败: {e}")
                self._is_initialized = False
                return False
    
    def _clear_pool(self):
        """清空连接池"""
        try:
            count = 0
            while not self.pool.empty():
                conn = self.pool.get_nowait()
                if conn:
                    try:
                        conn.close()
                        count += 1
                    except Exception as e:
                        logger.error(f"关闭数据库连接失败: {e}")
            if count > 0:
                logger.info(f"已关闭 {count} 个数据库连接")
        except Exception as e:
            logger.error(f"清空连接池失败: {e}")
        self._is_initialized = False
        
    def get_connection(self) -> Optional[sqlite3.Connection]:
        """获取一个连接"""
        if not self._is_initialized:
            return None
            
        try:
            # 从池中获取连接
            conn = self.pool.get(timeout=5)
            # 验证连接是否有效
            try:
                conn.execute("SELECT 1").fetchone()
                return conn
            except sqlite3.Error:
                # 如果连接失效，创建新连接
                try:
                    conn = sqlite3.connect(self.db_path, check_same_thread=False)
                    conn.execute("PRAGMA foreign_keys = ON")
                    conn.execute("PRAGMA journal_mode = WAL")
                    conn.execute("PRAGMA busy_timeout = 10000")
                    conn.execute("PRAGMA synchronous = NORMAL")
                    return conn
                except Exception as e:
                    logger.error(f"创建新数据库连接失败: {e}")
                    return None
        except queue.Empty:
            logger.warning("获取连接超时，连接池可能已满")
            return None
    
    def release_connection(self, conn):
        """释放连接回池中"""
        if conn and self._is_initialized:
            try:
                # 确保没有未完成的事务
                try:
                    conn.commit()  # 提交任何挂起的事务
                except sqlite3.Error as e:
                    logger.warning(f"释放连接时提交事务失败: {e}")
                    try:
                        conn.rollback()  # 如果提交失败，回滚
                    except:
                        pass
                
                # 返回池中
                self.pool.put(conn, timeout=5)
            except queue.Full:
                # 如果池已满，关闭连接
                try:
                    conn.close()
                    logger.info("连接池已满，关闭多余连接")
                except Exception as e:
                    logger.error(f"关闭多余连接失败: {e}")
                
    def close_all(self):
        """关闭所有连接"""
        with self.lock:
            self._clear_pool()
            
    @property
    def is_initialized(self):
        return self._is_initialized

# 创建全局连接池实例
db_pool = DatabasePool() 